import torch
import models
import numpy as np
from PIL import Image
import cv2
from data_cell import prepare_image_cv2
import os


resume = './ckpt_multi/crack_512.pth'
img_dir = r'H:\yuanbaoxi\ybx_gitee\RCF_EDGE\dataset\train\imgs'
label_dir = r'H:\yuanbaoxi\ybx_gitee\RCF_EDGE\dataset\train\edge'
result_path = './examples/0.png'


model = models.resnet101(pretrained=False).cuda()
model.eval()
#resume..
checkpoint = torch.load(resume)
model.load_state_dict(checkpoint)
names = os.listdir(img_dir)


for name in names:
    img_path = os.path.join(img_dir, name)
    name_s = name.split(".")
    name_l = name_s[0] + ".png"
    lb_file = os.path.join(label_dir, name_l)
    lb_uint8 = cv2.imread(lb_file)
    cv2.imshow("lb_uint8", cv2.resize(lb_uint8, (600, 600)))

    img_uint8 = cv2.imread(img_path)
    original_img = np.array(img_uint8, dtype=np.float32)
    cv2.imshow("img_uint8", cv2.resize(img_uint8, (600, 600)))
    h, w, _ = original_img.shape

    img = prepare_image_cv2(original_img)
    img = torch.from_numpy(img).unsqueeze(0).cuda()

    outs = model(img, (h, w))
    result = outs[-1].squeeze().detach().cpu().numpy()

    result = (result * 255).astype(np.uint8)
    cv2.imshow("result",cv2.resize(result, (600, 600)))
    cv2.waitKey(1000)
    # Image.fromarray(result).save(result_path)

